"""
Base Architecture for Multi-Task Learning
多任务学习架构基类
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Any
from ml_lib.core import Tensor
from ml_lib.nn.module import Module


class MTLArchitecture(Module, ABC):
    """多任务学习架构基类"""
    
    def __init__(self, input_dim: int, task_configs: Dict[str, Any]):
        """
        初始化MTL架构
        
        Args:
            input_dim: 输入维度
            task_configs: 任务配置字典
        """
        super().__init__()
        self.input_dim = input_dim
        self.task_configs = task_configs
        self.task_names = list(task_configs.keys())
        self.num_tasks = len(self.task_names)
        
    @abstractmethod
    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """
        前向传播
        
        Args:
            x: 输入张量
            
        Returns:
            每个任务的输出字典
        """
        pass
    
    @abstractmethod
    def get_shared_representation(self, x: Tensor) -> Tensor:
        """
        获取共享表示
        
        Args:
            x: 输入张量
            
        Returns:
            共享表示张量
        """
        pass
    
    def get_task_specific_params(self, task_name: str) -> List[Any]:
        """
        获取特定任务的参数
        
        Args:
            task_name: 任务名称
            
        Returns:
            任务特定参数列表
        """
        # 默认实现，子类可以重写
        return []
    
    def get_shared_params(self) -> List[Any]:
        """
        获取共享参数
        
        Returns:
            共享参数列表
        """
        # 默认实现，子类可以重写
        return [] 